from prophet import Prophet
from common_import import *
import mydraw
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score


def predict_quantity(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 只使用 'date' 和 'total_quantity' 进行预测
    df.rename(columns={"date": "ds", "total_quantity": "y"}, inplace=True)

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(
        country_name="CN"
    )  # 添加美国的节假日（根据需要选择国家）

    # 拟合模型
    model.fit(df)

    # 生成未来 7 天的日期
    future = model.make_future_dataframe(periods=7)

    # 进行预测
    forecast = model.predict(future)

    # 提取未来 7 天的预测值
    predicted_values = forecast[["ds", "yhat"]].tail(7)

    # 转换为结构化 NumPy 数组
    structured_array = np.array(
        list(
            zip(
                predicted_values["ds"].dt.strftime("%Y-%m-%d"), predicted_values["yhat"]
            )
        ),
        dtype=[("date", "U10"), ("total_quantity", "f4")],
    )

    return structured_array


def predict_sale_price(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 只使用 'date' 和 'total_quantity' 进行预测
    df.rename(columns={"date": "ds", "sale_price_index": "y"}, inplace=True)

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(
        country_name="CN"
    )  # 添加美国的节假日（根据需要选择国家）

    # 拟合模型
    model.fit(df)

    # 生成未来 7 天的日期
    future = model.make_future_dataframe(periods=7)

    # 进行预测
    forecast = model.predict(future)

    # 提取未来 7 天的预测值
    predicted_values = forecast[["ds", "yhat"]].tail(7)

    # 转换为结构化 NumPy 数组
    structured_array = np.array(
        list(
            zip(
                predicted_values["ds"].dt.strftime("%Y-%m-%d"), predicted_values["yhat"]
            )
        ),
        dtype=[("date", "U10"), ("sale_price", "f4")],
    )

    return structured_array


def predict_cost_price(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 只使用 'date' 和 'total_quantity' 进行预测
    df.rename(columns={"date": "ds", "cost_price_index": "y"}, inplace=True)

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(
        country_name="CN"
    )  # 添加中国的节假日（根据需要选择国家）

    # 拟合模型
    model.fit(df)

    # 生成未来 7 天的日期
    future = model.make_future_dataframe(periods=7)

    # 进行预测
    forecast = model.predict(future)

    # 提取未来 7 天的预测值
    predicted_values = forecast[["ds", "yhat"]].tail(7)

    # 转换为结构化 NumPy 数组
    structured_array = np.array(
        list(
            zip(
                predicted_values["ds"].dt.strftime("%Y-%m-%d"), predicted_values["yhat"]
            )
        ),
        dtype=[("date", "U10"), ("cost_price", "f4")],
    )

    return structured_array


def predict_all(data):
    result1 = predict_quantity(data)
    result2 = predict_sale_price(data)
    result3 = predict_cost_price(data)
    dates = result1["date"]
    merged_results = np.array(
        list(
            zip(
                dates,
                result1["total_quantity"],
                result2["sale_price"],  # 假设result2的结果是sale_price
                result3["cost_price"],  # 假设result3的结果是cost_price
            )
        ),
        dtype=[
            ("date", "U10"),
            ("total_quantity", "f4"),
            ("sale_price", "f4"),
            ("cost_price", "f4"),
        ],
    )

    return merged_results


def generate_6type_predict_csv():
    for i in range(1, 7):
        data = tool.get_np(f"problem_2/{i}.csv")
        result = predict_all(data)
        tool.get_csv(result, f"predict_2/{i}.csv")


def mydraw(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 只使用 'date' 和 'total_quantity' 进行预测
    df.rename(columns={"date": "ds", "total_quantity": "y"}, inplace=True)

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(country_name="CN")  # 添加中国的节假日

    # 拟合模型
    model.fit(df)

    # 生成未来 7 天的日期
    future = model.make_future_dataframe(periods=7)

    # 进行预测
    forecast = model.predict(future)

    # 画出分解图像
    fig = model.plot_components(forecast)
    for ax in fig.get_axes():
        ax.set_xlabel("")
    dir_path = "c:/Users/16028/Desktop/mcm/2023c/proj/picture"
    file_name = "prophet.png"
    import os

    if file_name:
        file_path = os.path.join(dir_path, file_name)
        plt.savefig(file_path, bbox_inches="tight", dpi=300)
    else:
        plt.show()


def analyze_seasonality_and_correlation(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 使用 'date' 和 'total_quantity' 进行 Prophet 预测
    df_prophet = df.rename(columns={"date": "ds", "total_quantity": "y"})

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(country_name="CN")  # 添加中国的节假日

    # 拟合模型
    model.fit(df_prophet)

    # 预测
    forecast = model.predict(df_prophet[["ds"]])

    # 从 total_quantity 中移除季节性、工作日和节假日效应
    print(forecast.columns)
    # df["deseasonalized_quantity"] = df["total_quantity"] - forecast["weekly"]
    df["deseasonalized_quantity"] = df["total_quantity"] - forecast["weekly"]
    # 对去除效应后的 total_quantity 与 sale_price_index 进行线性回归
    X = df["sale_price_index"].values.reshape(-1, 1)
    y = df["deseasonalized_quantity"].values

    model_lr = LinearRegression()
    model_lr.fit(X, y)
    y_pred = model_lr.predict(X)
    print("截距 (Intercept):", model_lr.intercept_)
    print("回归系数 (Coefficients):", model_lr.coef_)
    print("R² 值 (R² score):", r2_score(y, y_pred))

    # 绘制结果
    plt.figure(figsize=(10, 6))

    # 绘制原始数据和去除效应后的数据
    plt.subplot(2, 1, 1)
    plt.plot(df["date"], df["total_quantity"], label="原始数量")
    plt.plot(df["date"], df["deseasonalized_quantity"], label="去除季节性后的数量")
    plt.legend()
    plt.title("原始数量 vs 去除季节性后的数量")

    # 绘制线性回归结果
    plt.subplot(2, 1, 2)
    plt.scatter(df["sale_price_index"], df["deseasonalized_quantity"], label="数据点")
    plt.plot(df["sale_price_index"], y_pred, color="red", label="线性拟合")
    plt.xlabel("销售价格指数")
    plt.ylabel("去除季节性后的数量")
    plt.legend()
    plt.title(f"线性回归: R²={r2_score(y, y_pred):.4f}")

    plt.tight_layout()
    plt.show()

    return model_lr


def get_X(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 使用 'date' 和 'total_quantity' 进行 Prophet 预测
    df_prophet = df.rename(columns={"date": "ds", "sale_price_index": "y"})
    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(country_name="CN")  # 添加中国的节假日
    # 拟合模型
    model.fit(df_prophet)
    # 预测
    forecast = model.predict(df_prophet[["ds"]])
    # 从 total_quantity 中移除季节性、工作日和节假日效应
    print(forecast.columns)
    # df["deseasonalized_quantity"] = df["total_quantity"] - forecast["weekly"]
    df["deseasonalized_quantity"] = (
        df["sale_price_index"] - forecast["weekly"] - forecast["yearly"]
    )
    # 对去除效应后的 total_quantity 与 sale_price_index 进行线性回归
    X = df["deseasonalized_quantity"].values.reshape(-1, 1)
    return X


def get_y(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 使用 'date' 和 'total_quantity' 进行 Prophet 预测
    df_prophet = df.rename(columns={"date": "ds", "total_quantity": "y"})

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(country_name="CN")  # 添加中国的节假日

    # 拟合模型
    model.fit(df_prophet)

    # 预测
    forecast = model.predict(df_prophet[["ds"]])

    # 从 total_quantity 中移除季节性、工作日和节假日效应
    print(forecast.columns)
    # df["deseasonalized_quantity"] = df["total_quantity"] - forecast["weekly"]
    df["deseasonalized_quantity"] = (
        df["total_quantity"] - forecast["weekly"] - forecast["yearly"]
    )
    # 对去除效应后的 total_quantity 与 sale_price_index 进行线性回归
    y = df["deseasonalized_quantity"].values
    return y


import statsmodels.api as sm


def log_log_demand_model_show(a, b):
    a = array_1d_ravel = np.ravel(a)
    ln_quantities = np.log(a)
    ln_sale_price_indices = np.log(b)
    # 3. 构建线性回归模型
    X = ln_sale_price_indices
    X = sm.add_constant(X)  # 添加常数项
    y = ln_quantities
    model = sm.OLS(y, X).fit()

    # 4. 打印回归结果和参数
    print("回归方程:")
    print(
        f"ln(Quantity) = {model.params[0]:.4f} + {model.params[1]:.4f} * ln(Sale Price Index)"
    )
    print("\n回归模型的参数：")
    print(f"截距 (Intercept): {model.params[0]}")
    print(f"回归系数 (Coefficient for ln(Sale Price Index)): {model.params[1]}")
    print(f"R^2: {model.rsquared}")
    print(f"P-value: {model.pvalues[1]}")

    # 5. 可视化回归结果
    plt.figure(figsize=(10, 6))
    plt.scatter(ln_sale_price_indices, ln_quantities, label="Observed data")
    plt.plot(
        ln_sale_price_indices, model.fittedvalues, color="red", label="Fitted line"
    )
    plt.xlabel("ln(价格)")
    plt.ylabel("ln(需求)")
    plt.title("双对数模型")
    plt.legend()
    # mydraw.show_or_print("loglog.png")
    plt.show()

    return model.params[1], model.rsquared, model.pvalues[1]


def linear_analysis(X, y):
    model_lr = LinearRegression()
    model_lr.fit(X, y)
    y_pred = model_lr.predict(X)
    print("截距 (Intercept):", model_lr.intercept_)
    print("回归系数 (Coefficients):", model_lr.coef_)
    print("R² 值 (R² score):", r2_score(y, y_pred))
    # 绘制散点图和回归直线
    plt.figure(figsize=(8, 6))

    # 绘制数据点
    plt.scatter(X, y, color="blue", label="数据点")

    # 绘制回归直线
    plt.plot(X, y_pred, color="red", linewidth=2, label="回归直线")

    # 添加图表标题和标签
    plt.title("线性回归分析")
    plt.xlabel("自变量 (X)")
    plt.ylabel("因变量 (y)")
    plt.legend()

    # 显示图形
    plt.show()


def draw_two_line(data):
    # 将 NumPy 数组转换为 DataFrame
    df = pd.DataFrame(
        data, columns=["date", "total_quantity", "sale_price_index", "cost_price_index"]
    )
    df["date"] = pd.to_datetime(df["date"])

    # 使用 'date' 和 'total_quantity' 进行 Prophet 预测
    df_prophet = df.rename(columns={"date": "ds", "total_quantity": "y"})

    # 初始化 Prophet 模型，考虑季节性和工作日效应
    model = Prophet(
        weekly_seasonality=True, yearly_seasonality=True, daily_seasonality=False
    )
    model.add_country_holidays(country_name="CN")  # 添加中国的节假日

    # 拟合模型
    model.fit(df_prophet)

    # 进行预测
    forecast = model.predict(df_prophet[["ds"]])

    # 提取实际值和预测值
    df_prophet["yhat"] = forecast["yhat"]

    # 绘制实际值和拟合值的曲线图
    plt.figure(figsize=(12, 6))

    plt.plot(
        df_prophet["ds"][-100:],
        df_prophet["y"][-100:],
        label="实际值",
        color="blue",
    )
    plt.plot(
        df_prophet["ds"][-100:],
        df_prophet["yhat"][-100:],
        label="拟合值",
        color="red",
        linestyle="--",
    )

    plt.xlabel("日期")
    plt.ylabel("总量")
    plt.title("实际值 vs 拟合值")
    plt.legend()
    plt.grid(True)
    plt.show()


if __name__ == "__main__":
    # mydraw(tool.get_np("problem_2/1.csv"))
    data = tool.get_np("problem_2/1.csv")
    # log_log_demand_model_show(get_X(data), get_y(data))
    # linear_analysis(get_X(data), get_y(data))
    draw_two_line(data)
    pass
